package cn.jbolt.ai.core.fileSplit;

import cn.jbolt.ai.klb.FileState;
import cn.jbolt.ai.klb.service.AiFileService;
import cn.jbolt.util.cache.CacheUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

@Component
public class FileSplitCache {



    private static AiFileService aiFileService;
    @Autowired
    public void setAiFileService(AiFileService aiFileService) {
        FileSplitCache.aiFileService = aiFileService;
    }

    /**
     * 获取文件分割的计数缓存
     * @param fileId
     * @return
     */
    public static FileSplitCounter getFileSplitCounter(Long fileId) {
        if (fileId == null) {
            return null;
        }
        return CacheUtil.getObject("file_split_task_" + fileId);
    }

    /**
     * 是否有失败的块
     * @param fileId
     * @return
     */
    public static boolean hasFailed(Long fileId) {
        FileSplitCounter fileSplitCount = getFileSplitCounter(fileId);
        if (fileSplitCount == null) {
            return false;
        }
        return fileSplitCount.getFailCount() > 0;
    }


    /**
     * 添加成功计数
     * @param fileId
     */
    public static void addSuccessCount(Long fileId) {
        FileSplitCounter fileSplitCount = getFileSplitCounter(fileId);
        if (fileSplitCount == null) {
            return;
        }
        fileSplitCount.addSuccessCount();
        if (fileSplitCount.getTotalCount() == fileSplitCount.getSuccessCount()) {
            //更新文件状态为成功
            aiFileService.updateState(fileId, FileState.EMBEDDING_SUCCESS.getCode());
        }
    }

    /**
     * 添加失败计数
     * @param fileId
     */
    public static void addFailCount(Long fileId) {
        FileSplitCounter fileSplitCount = getFileSplitCounter(fileId);
        if (fileSplitCount == null) {
            return;
        }
        fileSplitCount.addFailCount();
        aiFileService.updateState(fileId, FileState.EMBEDDING_ERROR.getCode());
    }
    /**
     * 添加总计数
     * @param fileId
     */
    public static void addTotalCount(Long fileId) {
        FileSplitCounter fileSplitCount = getFileSplitCounter(fileId);
        if (fileSplitCount == null) {
            return;
        }
        fileSplitCount.addTotalCount();
    }




    /**
     * 设置文件分割的计数缓存
     * @param fileId
     */
    public static void addFileSplitCounter(Long fileId) {
        CacheUtil.putObject("file_split_task_" + fileId, new FileSplitCounter(), 6, TimeUnit.HOURS);
    }

}
